{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import symforce\n",
    "\n",
    "symforce.set_symbolic_api(\"sympy\")\n",
    "symforce.set_log_level(\"warning\")\n",
    "\n",
    "import symforce.symbolic as sf\n",
    "\n",
    "epsilon = 1e-9\n",
    "\n",
    "import numpy as np\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "frame_count = 100\n",
    "loops = 7\n",
    "slider_steps = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Functions for Generating Figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tangent_vectors(a: sf.Unit3):\n",
    "    basis = a.basis(epsilon)\n",
    "    return basis[:, 0].to_numpy().flatten(), basis[:, 1].to_numpy().flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate an animation that rotates around the unit sphere.\n",
    "\n",
    "\n",
    "def create_frames(num_frames):\n",
    "    frames = []\n",
    "    for i in range(num_frames):\n",
    "        v = get_unit_vector(i)\n",
    "        v_draw = v.to_unit_vector().to_numpy().flatten()\n",
    "        t1, t2 = get_tangent_vectors(v)\n",
    "        frame_data = [\n",
    "            create_unit_sphere(),\n",
    "            create_vector_trace(np.zeros(3), v_draw, \"Unit3\"),\n",
    "            create_vector_trace(v_draw, t1, \"Tangent X\"),\n",
    "            create_vector_trace(v_draw, t2, \"Tangent Y\"),\n",
    "        ]\n",
    "        frames.append(go.Frame(data=frame_data, name=str(i)))\n",
    "    return frames\n",
    "\n",
    "\n",
    "def get_unit_vector(i):\n",
    "    phi = np.pi * i / frame_count\n",
    "    theta = 2 * loops * np.pi * i / frame_count\n",
    "    return sf.Unit3(sf.V3(np.cos(phi), np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta)))\n",
    "\n",
    "\n",
    "def create_vector_trace(a: np.array, b: np.array, name: str = None, color: str = None):\n",
    "    return go.Scatter3d(\n",
    "        x=[a[0], a[0] + b[0]],\n",
    "        y=[a[1], a[1] + b[1]],\n",
    "        z=[a[2], a[2] + b[2]],\n",
    "        mode=\"lines+markers\",\n",
    "        marker=dict(size=3),\n",
    "        line=dict(width=3, color=color),\n",
    "        name=name,\n",
    "    )\n",
    "\n",
    "\n",
    "def create_unit_sphere():\n",
    "    u = np.linspace(0, 2 * np.pi, 20)\n",
    "    v = np.linspace(0, np.pi, 20)\n",
    "\n",
    "    x = np.outer(np.cos(u), np.sin(v))\n",
    "\n",
    "    # Plot the sphere with a slight grey gradient. and transparent.\n",
    "    return go.Surface(\n",
    "        x=x,\n",
    "        y=np.outer(np.sin(u), np.sin(v)),\n",
    "        z=np.outer(np.ones(np.size(u)), np.cos(v)),\n",
    "        colorscale=[[0, \"rgb(220,220,220)\"], [1, \"rgb(220,220,220)\"]],\n",
    "        opacity=0.5,\n",
    "        showscale=False,\n",
    "    )\n",
    "\n",
    "\n",
    "def build_animation():\n",
    "    initial_v = get_unit_vector(0)\n",
    "    t1, t2 = get_tangent_vectors(initial_v)\n",
    "\n",
    "    v_draw = initial_v.to_unit_vector().to_numpy().flatten()\n",
    "\n",
    "    frames = create_frames(frame_count)\n",
    "\n",
    "    # Define slider steps.\n",
    "    slider_steps = []\n",
    "    for i, frame in enumerate(frames):\n",
    "        step = {\n",
    "            \"method\": \"animate\",\n",
    "            \"label\": f\"{i}\",\n",
    "            \"args\": [\n",
    "                [frame.name],\n",
    "                {\n",
    "                    \"mode\": \"immediate\",\n",
    "                    \"frame\": {\"redraw\": True, \"duration\": 0},\n",
    "                    \"transition\": {\"duration\": 0},\n",
    "                },\n",
    "            ],\n",
    "        }\n",
    "        slider_steps.append(step)\n",
    "\n",
    "    fig = go.Figure(\n",
    "        data=[\n",
    "            create_unit_sphere(),\n",
    "            create_vector_trace(np.zeros(3), v_draw, \"Unit3\", \"black\"),\n",
    "            create_vector_trace(v_draw, t1, \"Tangent X\", \"blue\"),\n",
    "            create_vector_trace(v_draw, t2, \"Tangent Y\", \"red\"),\n",
    "        ],\n",
    "        layout=go.Layout(\n",
    "            scene=dict(\n",
    "                xaxis=dict(range=[-1.5, 1.5]),\n",
    "                yaxis=dict(range=[-1.5, 1.5]),\n",
    "                zaxis=dict(range=[-1.5, 1.5]),\n",
    "                aspectmode=\"cube\",\n",
    "            ),\n",
    "            sliders=[\n",
    "                {\n",
    "                    \"active\": 0,\n",
    "                    \"currentvalue\": {\"prefix\": \"Frame: \"},\n",
    "                    \"pad\": {\"t\": 50},\n",
    "                    \"steps\": slider_steps,\n",
    "                }\n",
    "            ],\n",
    "        ),\n",
    "        frames=frames,\n",
    "    )\n",
    "\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a figure that just shows the the tangent x / y vectors across the unit sphere.\n",
    "\n",
    "\n",
    "def compute_unit_vectors(n):\n",
    "    \"\"\"\n",
    "    Returns n unit vectors evenly distributed on a sphere.\n",
    "    \"\"\"\n",
    "    points = []\n",
    "    golden_angle = np.pi * (3 - np.sqrt(5))\n",
    "    for i in range(n):\n",
    "        y = 1 - (i / float(n - 1)) * 2  # y goes from 1 to -1\n",
    "        radius = np.sqrt(1 - y * y)\n",
    "        theta = golden_angle * i\n",
    "        x = np.cos(theta) * radius\n",
    "        z = np.sin(theta) * radius\n",
    "        points.append(sf.Unit3(sf.V3(x, y, z)))\n",
    "    return points\n",
    "\n",
    "\n",
    "def group_vector_traces(vectors, color, name):\n",
    "    \"\"\"\n",
    "    Given vectors as list of (start, end) pairs (each a 3-array),\n",
    "    returns a single trace that draws them as separate segments.\n",
    "    \"\"\"\n",
    "    xs, ys, zs = [], [], []\n",
    "    # For each segment, append start and end coords, then a gap (None).\n",
    "    for start, end in vectors:\n",
    "        # Start point.\n",
    "        xs.append(start[0])\n",
    "        ys.append(start[1])\n",
    "        zs.append(start[2])\n",
    "        # End point.\n",
    "        xs.append(end[0])\n",
    "        ys.append(end[1])\n",
    "        zs.append(end[2])\n",
    "        # Insert gap.\n",
    "        xs.append(None)\n",
    "        ys.append(None)\n",
    "        zs.append(None)\n",
    "    return go.Scatter3d(\n",
    "        x=xs,\n",
    "        y=ys,\n",
    "        z=zs,\n",
    "        mode=\"lines+markers\",\n",
    "        marker=dict(size=3),\n",
    "        line=dict(width=3, color=color),\n",
    "        name=name,\n",
    "    )\n",
    "\n",
    "\n",
    "def build_tangent_figure():\n",
    "    unit_vectors = compute_unit_vectors(frame_count)\n",
    "\n",
    "    data = [create_unit_sphere()]\n",
    "\n",
    "    vectors_x = []\n",
    "    vectors_y = []\n",
    "    for v in unit_vectors:\n",
    "        v_draw = v.to_unit_vector().to_numpy().flatten()\n",
    "        t1, t2 = get_tangent_vectors(v)\n",
    "        vectors_x.append((v_draw, v_draw + t1 * 0.1))\n",
    "        vectors_y.append((v_draw, v_draw + t2 * 0.1))\n",
    "\n",
    "    data.append(group_vector_traces(vectors_x, \"blue\", \"Tangent X\"))\n",
    "    data.append(group_vector_traces(vectors_y, \"red\", \"Tangent Y\"))\n",
    "\n",
    "    # Adjust the size to help visualize the tangent vectors.\n",
    "    fig = go.Figure(\n",
    "        data=tuple(data),\n",
    "        layout=go.Layout(\n",
    "            scene=dict(\n",
    "                xaxis=dict(range=[-1.5, 1.5]),\n",
    "                yaxis=dict(range=[-1.5, 1.5]),\n",
    "                zaxis=dict(range=[-1.5, 1.5]),\n",
    "                aspectmode=\"cube\",\n",
    "            ),\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    fig.show()\n",
    "\n",
    "\n",
    "def build_tangent_figure_singularity():\n",
    "    # Only select unit vectors around the singularity direction at (1, 0, 0).\n",
    "    unit_vectors = []\n",
    "    for i in range(frame_count):\n",
    "        golden_angle = np.pi * (3 - np.sqrt(5))\n",
    "        phi = 0.05 * np.pi * i / frame_count\n",
    "        theta = 2 * i * 100 * golden_angle / frame_count\n",
    "        unit_vectors.append(\n",
    "            sf.Unit3(sf.V3(np.cos(phi), np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta)))\n",
    "        )\n",
    "    # unit_vectors = compute_unit_vectors(frame_count)\n",
    "    data = [create_unit_sphere()]\n",
    "    vectors_x = []\n",
    "    vectors_y = []\n",
    "    for v in unit_vectors:\n",
    "        v_draw = v.to_unit_vector().to_numpy().flatten()\n",
    "        t1, t2 = get_tangent_vectors(v)\n",
    "        vectors_x.append((v_draw, v_draw + t1 * 0.01))\n",
    "        vectors_y.append((v_draw, v_draw + t2 * 0.01))\n",
    "\n",
    "    data.append(group_vector_traces(vectors_x, \"blue\", \"Tangent X\"))\n",
    "    data.append(group_vector_traces(vectors_y, \"red\", \"Tangent Y\"))\n",
    "\n",
    "    # Adjust the size to help visualize the tangent vectors.\n",
    "    fig = go.Figure(\n",
    "        data=tuple(data),\n",
    "        layout=go.Layout(\n",
    "            scene=dict(\n",
    "                xaxis=dict(range=[-1.5, 1.5]),\n",
    "                yaxis=dict(range=[-1.5, 1.5]),\n",
    "                zaxis=dict(range=[-1.5, 1.5]),\n",
    "                aspectmode=\"cube\",\n",
    "            ),\n",
    "        ),\n",
    "    )\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display Unit3 Tangent Space Sweep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "build_animation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display Unit3 Tangent X / Y Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "build_tangent_figure()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display Unit3 Tangent X / Y Space -- Around Singularity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "build_tangent_figure_singularity()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
